import streamlit as st
from PIL import Image

from utils.Detect import *
from utils.LoadModel import *

mission = '舰船检测'
st.header(mission)
print('-' * 20, mission, '-' * 20)
model = load_sea_ships_model()

func = st.sidebar.selectbox(
    '选择数据源',
    ['图片检测', '摄像头']
)

st.logo(Config.logo)

# 图片与模型
if func == '图片检测':
    uploaded_img = st.file_uploader("上传照片进行检测", type=["jpg", "jpeg", "png"])
    col1, col2 = st.columns(2)

    if uploaded_img is not None:
        img = Image.open(uploaded_img)
        with col1:
            st.image(img)

        results = detect(model, img)
        img_draw = draw_img(model, results, img)

        path = Config.hisPic + f'{time.time()}' + '.jpg'
        img_draw.save(path)

        save_data(mission, path)

        with col2:
            st.image(img_draw)

elif func == '摄像头':
    st.markdown('<span style="color: red; font-size: 50px">该项目不支持摄像头检测</span>', unsafe_allow_html=True)
